package se.webbzon.oschi01.aes;

public class RijndaelCipher {
	
	// The S-box used for this Rijndael cipher
	private final RijndaelSBox sBox;
	
	// An array of all the subkeys used for this Rijndael cipher
	private RijndaelMatrix[] subKeys;
	
	/*============================================================
	Constructors
	============================================================*/
	
	/** Creates a new Rijndael cipher instance. **/
	public RijndaelCipher() {
		this.sBox = new RijndaelSBox();
	}
	
	/** Creates a new Rijndael cipher instance with the provided
	 * S-box. **/
	public RijndaelCipher(RijndaelSBox sBox) {
		this.sBox = sBox;
	}
	
	/*============================================================
	Public Methods
	============================================================*/
	
	/** Sets the key used for this Rijndael cipher. **/
	public void setKey(RijndaelMatrix key) {
		if (key == null) {
			subKeys = null;
		} else {
			subKeys = new RijndaelMatrix[11];
			subKeys[0] = key;
			for (byte i = 1; i < 11; i++)
				subKeys[i] = subKeys[i-1].createSubKey(sBox, i);
		}
	}
	
	/** Returns the key used for this Rijndael cipher. **/
	public RijndaelMatrix getKey() {
		if (subKeys == null)
			return null;
		else
			return subKeys[0];
	}
	
	/** Encrypts the provided message by using the Rijndael
	 * cipher algorithm with the current key. **/
	public void encrypt(RijndaelMatrix[] msg) {
		if (subKeys == null)
			throw new IllegalStateException("No key has been selected for this Rijndael cipher.");
		encrypt(msg,subKeys);
	}
	
	/** Decrypts the provided message by using the Rijndael
	 * cipher algorithm with the current key. **/
	public void decrypt(RijndaelMatrix[] msg) {
		if (subKeys == null)
			throw new IllegalStateException("No key has been selected for this Rijndael cipher.");
		decrypt(msg,subKeys);
	}
	
	/** Encrypts the provided message by using the Rijndael
	 * cipher algorithm with the given key. **/
	public void encrypt(RijndaelMatrix[] msg, RijndaelMatrix key) {
		RijndaelMatrix[] subKeys = new RijndaelMatrix[11];
		subKeys[0] = key;
		for (byte i = 1; i < 11; i++)
			subKeys[i] = subKeys[i-1].createSubKey(sBox, i);
		encrypt(msg,subKeys);
	}
	
	/** Encrypts the provided message by using the Rijndael
	 * cipher algorithm with the given key. **/
	public void decrypt(RijndaelMatrix[] msg, RijndaelMatrix key) {
		RijndaelMatrix[] subKeys = new RijndaelMatrix[11];
		subKeys[0] = key;
		for (byte i = 1; i < 11; i++)
			subKeys[i] = subKeys[i-1].createSubKey(sBox, i);
		decrypt(msg,subKeys);
	}
	
	/*============================================================
	Private Methods
	============================================================*/
	
	/** Encrypts the given message by using the given set of subkeys. **/
	private void encrypt(RijndaelMatrix[] msg, RijndaelMatrix[] subKeys) {
		for (RijndaelMatrix m : msg) {
			m.addRoundKey(subKeys[0]);
			for (int i = 1; i < 10; i++) {
				m.subBytes(sBox);
				m.shiftRows();
				m.mixColumns();
				m.addRoundKey(subKeys[i]);
			}
			m.subBytes(sBox);
			m.shiftRows();
			m.addRoundKey(subKeys[10]);
		}
	}
	
	/** Decrypts the given message by using the given set of subkeys. **/
	private void decrypt(RijndaelMatrix[] msg, RijndaelMatrix[] subKeys) {
		for (RijndaelMatrix m : msg) {
			m.addRoundKey(subKeys[10]);
			m.invShiftRows();
			m.invSubBytes(sBox);
			for (int i = 9; i > 0; i--) {
				m.addRoundKey(subKeys[i]);
				m.invMixColumns();
				m.invShiftRows();
				m.invSubBytes(sBox);
			}
			m.addRoundKey(subKeys[0]);
		}
	}

}
